package com.solidsnake.scala.http

import java.nio.charset.Charset
import java.net.{HttpURLConnection, URLConnection, URL}

import java.io.{ByteArrayOutputStream, PrintWriter, InputStream}

import scala.collection.JavaConverters._


object HttpUtils {

  type ResponseInfo = (Map[String, String], Map[String, String], Array[Byte])

  implicit def setMap2Connection(conn: HttpURLConnection) = {
    new {
      def setRequestHeader(header: Map[String, String]) = {
        header.foreach {
          case (k, v) => conn.setRequestProperty(k, v)
        }
        conn
      }
    }
  }


  def getText = request("get", "") _

  def getImage = request("get", Array[Byte]()) _

  def postText = request("post", "") _

  def postImage = request("post", Array[Byte]()) _


  /**
   *
   * @param urlStr
   * @param params
   * @param requestHandler
   * @return
   */
  def get(urlStr: String, params: Map[String, String], requestHandler: Map[String, String] ): (Map[String, String], Map[String, String], Array[Byte]) = {
    var in: InputStream = null
    var conn: HttpURLConnection = null

    val url = new URL(urlStr + "?" + this.paramMapToString(params))
    conn = url.openConnection().asInstanceOf[HttpURLConnection]
    conn.setRequestHeader(requestHandler)
    conn.setInstanceFollowRedirects(false)
    conn.connect()
    try {
      in = conn.getInputStream()
      val responseHandler = this.readHeader(conn)
      val body = this.readBody(in, responseHandler)
      return Tuple3(requestHandler, responseHandler, body)
    } finally {
      if (in != null) in.close()
    }
  }


  def post(urlStr: String, params: Map[String, String], requestHandler: Map[String, String]): (Map[String, String], Map[String, String], Array[Byte]) = {

    var conn: HttpURLConnection = null
    var out: PrintWriter = null
    var in: InputStream = null

    val url = new URL(urlStr)
    conn = url.openConnection().asInstanceOf[HttpURLConnection]
    conn.setRequestProperty("accept", "*/*")
    conn.setRequestProperty("connection", "Keep-Alive")
    conn.setRequestProperty("user-agent",
      "Mozilla/5.0 (compatible; MSIE 9.0; Windows NT 6.1; Trident/5.0)")
    conn.setInstanceFollowRedirects(false);
    conn.setRequestHeader(requestHandler)
    conn.setDoOutput(true)
    conn.setDoInput(true)

    try {
      out = new PrintWriter(conn.getOutputStream())
      out.print(this.paramMapToString(params))
      out.flush()

      in = conn.getInputStream()
      val responseHandler = this.readHeader(conn)
      val body = this.readBody(in, responseHandler)
      return Tuple3(requestHandler, responseHandler , body)
    } finally {
      if (out != null)
        out.close()
      if (in != null)
        in.close()
    }
  }


  def request[T](method: String, dataType: T)
                (url: String, params: Map[String, String],
                 charset: Charset = Charset.forName("UTF-8"),
                 header: Map[String, String] = Map(),
                 redirectCookie: Boolean = false): (Map[String,String], Map[String, String], T) = {

    val requestFuc = method match {
      case "get" => get(url, params, _: Map[String, String])
      case "post" => post(url, params, _: Map[String, String])
    }
    val responseData = redirectCookie match {
      case false => requestFuc(header)
      case true => {
        var flag = false
        var result: ResponseInfo = null
        var cookie = header.getOrElse(HttpHandlerField.COOKIE, "")
        while (!flag) {
          println("请求开始Cookie: " + cookie)
          val newHeader = header + (HttpHandlerField.COOKIE -> cookie)
          result = requestFuc(newHeader)

          val setCookie = result._2.getOrElse(HttpHandlerField.SET_COOKIE, "")
          if(!setCookie.isEmpty && cookie.indexOf(setCookie) == -1) {
            val oldCookie = newHeader.getOrElse(HttpHandlerField.COOKIE, "")
            cookie = oldCookie + ";" +result._2(HttpHandlerField.SET_COOKIE)
            println("获得COOKIE: " + cookie)
          } else {
            flag = true
          }
        }
        result
      }
    }

    dataType match {
      case t: String => Tuple3(responseData._1,responseData._2, new String(responseData._3, charset).asInstanceOf[T])
      case t: Array[Byte] => Tuple3(responseData._1,responseData._2, responseData._3.asInstanceOf[T])
    }
  }

  private def paramMapToString(params: Map[String, String]): String = {
    val paramStr = ("" /: params) {
      case (a, (k, v)) => a + "&" + k + "=" + v
    }
    if (paramStr.length > 0) {
      paramStr.substring(1)
    } else {
      paramStr
    }
  }


  private def readBody(in: InputStream, header: Map[String, String]) = {
    val size = header.getOrElse(HttpHandlerField.CONTENT_LENGTH, "1024")
    val bytes = new Array[Byte](Integer.parseInt(size));
    val byteOutput = new ByteArrayOutputStream();

    var length = in.read(bytes);
    while (length != -1) {
      byteOutput.write(bytes, 0, length);
      length = in.read(bytes)
    }
    byteOutput.toByteArray();
  }


  private def readHeader(conn: URLConnection): Map[String, String] = {
    var result = Map[String, String]()
    for ((k, v) <- conn.getHeaderFields.asScala) {
      result += (k -> ("" /: v.asScala)(_ + _))
    }
    result
  }


  def main(s: Array[String]): Unit = {
    //http://app02.szaic.gov.cn/WERMIS2.WebUI/Default.aspx
    val header = Map[String, String]()
    val params = Map("hfIp" -> "119.166.39.25")
    val result = HttpUtils.postText("http://app02.szaic.gov.cn/WERMIS2.WebUI/Default.aspx", params, Charset.forName("UTF-8"), header, false)
    val results = StringBuilder.newBuilder
    for ((k, v) <- result._1) {
      if(k != null) {
        results append k + ":"
      }
      v.foreach(value => {
        results append value
      })
      results append "\n"
    }
    println(results)

    println(new String(result._3))
  }

}
